import logging
from contextlib import nullcontext
from dataclasses import asdict
from pprint import pformat

import re
import torch

from lerobot.common.policies.factory import make_policy
from lerobot.common.policies.pretrained import PreTrainedPolicy
from lerobot.configs import parser
from lerobot.configs.eval import EvalPipelineConfig
from lerobot.common.utils.utils import get_safe_torch_device

from viewer_aloha_env import ViewAlohaEnv


def camel_to_snake(text):
    return re.sub(r'(?<!^)(?=[A-Z])', '_', text).lower()


def rollout(env: ViewAlohaEnv, policy: PreTrainedPolicy):
    policy.eval()

    policy.reset()
    env.set_policy(policy)

    env.reset()
    env.view()


@parser.wrap()
def main(cfg: EvalPipelineConfig):
    logging.info(pformat(asdict(cfg)))

    device = get_safe_torch_device(cfg.policy.device, log=True)

    policy = make_policy(cfg=cfg.policy, env_cfg=cfg.env)
    policy.eval()

    env = ViewAlohaEnv(task=camel_to_snake(cfg.env.task), obs_type=cfg.env.obs_type)

    with torch.no_grad(), torch.autocast(device_type=device.type) if cfg.policy.use_amp else nullcontext():
        rollout(env, policy)


if __name__ == '__main__':
    main()
